% This function takes estimates and imputes missing network and A


function [G_ij_imp, A_imp] = Fn_ImputeNetwork_Mar2022(G_ij, Est, Cons, net_tol, DELTA)

G_ij_imp = G_ij;
G_ji_imp = Cons.Flip_ij*G_ij_imp;

% Impute A
A_imp = Est.A_hat.*(ones(Cons.n,1)-Cons.A_MISS) + randn(Cons.n,1).*Cons.A_MISS;
A_j_imp = Cons.I_j*A_imp;
A_i_imp = Cons.I_i*A_imp;

[V, e] = eig(Est.V_C);
C_ij_imp_raw = randn(size(Cons.I_j,1),2)*(V*sqrt(e))';
C_ij_imp = (ones(size(Cons.J_greater,1),1) - Cons.J_greater).*C_ij_imp_raw(:,1) + Cons.J_greater.*(Cons.Flip_ij*C_ij_imp_raw(:,2));
%cov(C_ij_imp, Cons.Flip_ij*C_ij_imp)

% Impute M and l_M_imp
X_M = [ones(size(Cons.X_raw,1),1), Cons.X_raw, Cons.P, bsxfun(@times,Cons.X_raw,Cons.P)];

l_M_imp = X_M*Est.beta_M_hat + sqrt(Est.sigma2_M_hat)*randn(Cons.n,1);
l_M_imp = (ones(Cons.n,1) - Cons.A_MISS).*Est.l_M_hat + Cons.A_MISS.*l_M_imp;


% Impute Networks
B_ij_imp = [Cons.X_j, A_j_imp, Cons.X_i.*Cons.X_j, Cons.X_i.*(A_j_imp*ones(1,Cons.dimX)), A_i_imp*(ones(1,Cons.dimX)).*Cons.X_j, A_i_imp.*A_j_imp];
Lambda_imp = 0;

diff = 1000;
while diff > net_tol

    Sum_B = Cons.I_i'*exp([G_ji_imp, B_ij_imp]*Est.BGD_hat + DELTA*Cons.NETMISS);
    Lambda_imp_new = log(Sum_B) - l_M_imp;

    G_ij_new = [G_ji_imp, B_ij_imp]*Est.BGD_hat + DELTA*Cons.NETMISS - Cons.I_i*Lambda_imp_new - C_ij_imp;

    G_ij_new =  G_ij_imp.*(ones(size(Cons.NETMISS,1),1) - Cons.NETMISS) + G_ij_new.*Cons.NETMISS;

    diff_new = sqrt((Lambda_imp_new - Lambda_imp)'*(Lambda_imp_new - Lambda_imp));


    G_ij_imp = G_ij_new;
    G_ji_imp = Cons.Flip_ij*G_ij_imp;
    Lambda_imp = Lambda_imp_new;
    diff = diff_new;

end

end